import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class Lr(nn.Module):
    def __init__(self, idim):
        super(Lr, self).__init__()
        self.lin = nn.Linear(idim, 1)

    def forward(self, x):
        out = self.lin(x)
        return out


x_values = [i for i in range(1, 10)]
x_train = np.array(x_values, dtype=np.float32)  # 构造numpy
x_train = x_train.reshape(-1, 1)  # 为防止出错，一般情况下将数据转换为矩阵格式，此处为列向量矩阵
x_train =torch.from_numpy(x_train)

y_values = [2 * i + 1 for i in x_values]  # 假设y=2*x+1
y_train = np.array(y_values, dtype=np.float32)
y_train = y_train.reshape(-1, 1)
y_train =torch.from_numpy(y_train)

print (x_train.shape,y_train.shape)
loss=nn.MSELoss()
mod=Lr(1)
opt=torch.optim.SGD(mod.parameters(),lr=0.01)
for e in range(100):
    out=mod(x_train)
    l=loss(out,y_train)
    l.backward()
    opt.step()
    opt.zero_grad()
    print (l)
mod.eval()
print (111,mod(x_train))


for k, v in mod.state_dict().items():
    print (k,v)
